// 通过缩减Tensor ShapeInfo维度，优化栈空间
#define K_MAX_SHAPE_DIM 0

#include "kernel_operator.h"

using namespace AscendC;

#define ll long long

constexpr int32_t BUFFER_NUM = 1;
constexpr int32_t TQUE_NUM = 1;
constexpr int32_t BUFFER_NUM_SP = 1;
constexpr int32_t BUFFER_NUM_CS3 = 2;

template<typename T> class BitWiseLeftShift {
private:
    TPipe* pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<T> x1Gm;
    GlobalTensor<T> x2Gm;
    GlobalTensor<T> yGm;
    int blockIdx;
    ll offset;
    ll offset1;
    ll offset2;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int calcLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;

public:
    __aicore__ inline BitWiseLeftShift() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        // this->totalLength1 = totalLength1;
        // this->totalLength2 = totalLength2;
        // this->totalLength = totalLength;
        // this->dimNum1 = dimNum1;
        // this->dimNum2 = dimNum2;
        // this->dimNum = 3;
        // this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < 3; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
            if (dim[i] != 1) {
                this->dimNum = i + 1;
            }
        }
        // 打印每个dimnum和dim
        // printf("dimNum1: %d, dimNum2: %d, dimNum: %d\n", dimNum1, dimNum2, dimNum);
        // for (int i = 0; i < dimNum; i++) {
        //     printf("dim[%d]: %d, dim1[%d]: %d, dim2[%d]: %d\n", i, dim[i], i, dim1[i], i, dim2[i]);
        // }
        
        blockIdx = GetBlockIdx();
        calcLength = dim[dimNum - 1];
        offset1 = blockIdx * calcLength;
        // 计算indices（每一维的索引）
        indices[dimNum - 1] = 0;
        int tmp = blockIdx;
        for (int i = dimNum - 2; i >= 0; i--) {
            indices[i] = tmp % dim[i];
            tmp /= dim[i];
        }
        // printf("blockIdx: %d, indices: ", blockIdx);
        // for (int i = 0; i < dimNum; i++) {
        //     printf("%d ", indices[i]);
        // }
        // printf("\n");

        // 计算x2索引
        offset2 = 0;
        ll stride2 = 1;
        for (int i = dimNum - 1; i >= 0; i--) {
            int idx = indices[i] % dim2[i]; // 处理广播
            offset2 += idx * stride2;
            stride2 *= dim2[i];
        }
        // printf("offset1: %d, offset2: %d\n", offset1, offset2);
        

        x1Gm.SetGlobalBuffer((__gm__ T *)x1 + offset1, calcLength * sizeof(T));
        x2Gm.SetGlobalBuffer((__gm__ T *)x2 + offset2, dim2[dimNum - 1] * sizeof(T));
        yGm.SetGlobalBuffer((__gm__ T *)y + offset1, calcLength * sizeof(T));

        pipe->InitBuffer(inQueueX1, BUFFER_NUM, calcLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM, calcLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM, calcLength * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(calcLength * sizeof(T)), 
            0,
            0,
            0
        };
        DataCopyExtParams copyParams2 = {
            (uint16_t)1, 
            (uint32_t)(1 * sizeof(T)), 
            0,
            0,
            0
        };
        if (dim2[dimNum - 1] == 1) {
            // 处理x2为1的情况
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm, copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            x1Local = inQueueX1.DeQue<T>();
            DataCopyPad(x2Local, x2Gm, copyParams2, padParams);
            inQueueX2.EnQue(x2Local);
            x2Local = inQueueX2.DeQue<T>();
            ll x2 = x2Local(0);
            inQueueX2.FreeTensor(x2Local);
            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
            if constexpr (std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>) {
                ShiftLeft(yLocal, x1Local, (T)x2, calcLength);
            } else {
                for (int i = 0; i < calcLength; i++) {
                    ll x1 = x1Local(i);
                    ll y = x1 << x2;
                    yLocal(i) = y;
                }
            }
            inQueueX1.FreeTensor(x1Local);
            outQueueY.EnQue(yLocal);
            yLocal = outQueueY.DeQue<T>();
            DataCopyPad(yGm, yLocal, copyParams);
            outQueueY.FreeTensor(yLocal);
            return;
        }
        // 处理x2为多维的情况
        LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x1Local, x1Gm, copyParams, padParams);
        DataCopyPad(x2Local, x2Gm, copyParams, padParams);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
        x1Local = inQueueX1.DeQue<T>();
        x2Local = inQueueX2.DeQue<T>();
        LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
        if constexpr (std::is_same_v<T, int16_t>) {
            ShiftLeft(x2Local, x2Local, (int16_t)10, calcLength);
            Adds(x2Local, x2Local, (int16_t)0b0011110000000000, calcLength);
            Cast(x2Local, x2Local.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            Mul(yLocal, x1Local, x2Local, calcLength);
        } else if constexpr (std::is_same_v<T, int32_t>) {
            ShiftLeft(x2Local, x2Local, (int32_t)23, calcLength);
            Adds(x2Local, x2Local, (int32_t)0b00111111100000000000000000000000, calcLength);
            Cast(x2Local, x2Local.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);
            Mul(yLocal, x1Local, x2Local, calcLength);
        } else {
            for (int i = 0; i < calcLength; i++) {
                ll x1 = x1Local(i);
                ll x2 = x2Local(i);
                ll y = x1 << x2;
                yLocal(i) = y;
            }
        }
        outQueueY.EnQue(yLocal);
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
        yLocal = outQueueY.DeQue<T>();
        DataCopyPad(yGm, yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }
};


template<typename T> class BitWiseLeftShift_SP {
private:
    TPipe* pipe;
    TQue<TPosition::VECIN, TQUE_NUM> inQueueX1, inQueueX2;
    TQue<TPosition::VECOUT, TQUE_NUM> outQueueY;
    TQue<TPosition::VECCALC, TQUE_NUM> calcQueueX1, calcQueueX2;
    GlobalTensor<T> x1Gm, x2Gm, yGm;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;
    int blockNum;
    int blockIdx;
    int blockLength;
    int tailBlockNum;
    int tileLength;
    int tileNum;

public:
    __aicore__ inline BitWiseLeftShift_SP() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        // this->totalLength1 = totalLength1;
        // this->totalLength2 = totalLength2;
        // this->totalLength = totalLength;
        // this->dimNum1 = dimNum1;
        // this->dimNum2 = dimNum2;
        // this->dimNum = 3;
        // this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < 3; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
            if (dim[i] != 1) {
                this->dimNum = i + 1;
            }
        }
        
        blockNum = GetBlockNum();
        blockIdx = GetBlockIdx();
        // 先给每个核分配要计算的长度
        blockLength = (totalLength2 + blockNum - 1) / blockNum;
        // int tailBlockNum = totalLength % blockNum;
        // this->blockLength = blockLength;
        // this->tailBlockNum = tailBlockNum;
        // 每个循环里计算的长度
        int bufferNum = 1;
        if constexpr (std::is_same_v<T, int8_t>) {
            tileLength = 4096;
        } else if constexpr (std::is_same_v<T, int16_t>) {
            tileLength = 8192;
        } else if constexpr (std::is_same_v<T, int32_t>) {
            if (totalLength >= 10000000) {
                tileLength = 8192;
            } else {
                tileLength = 4096;
                bufferNum = 2;
            }
        } else {
            tileLength = 1;
        }
        if (tileLength > totalLength) {
            tileLength = totalLength;
        }
        // 向上对齐到256字节
        // tileLength = (tileLength * sizeof(T) + 255) / 256 * 256 / sizeof(T);

        // printf("blockIdx = %d, blockLength = %d, tailBlockNum = %d\n", blockIdx, blockLength, tailBlockNum);
        // printf("alreadyLength = %d, tileLength = %d, tileNum = %d\n", alreadyLength, tileLength, tileNum);

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, totalLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, totalLength);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);

        pipe->InitBuffer(inQueueX1, bufferNum, tileLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, bufferNum, tileLength * sizeof(T));
        pipe->InitBuffer(outQueueY, bufferNum, tileLength * sizeof(T));
        if constexpr (std::is_same_v<T, int8_t>) {
            pipe->InitBuffer(calcQueueX1, bufferNum, tileLength * sizeof(int16_t));
            pipe->InitBuffer(calcQueueX2, bufferNum, tileLength * sizeof(int16_t));
        }
    }
    __aicore__ inline void Process()
    {
        #define min(x, y) ((x) < (y) ? (x) : (y))
        for (size_t offset = blockIdx * tileLength; offset < totalLength; offset += blockNum * tileLength) {
            int calcLength = min(tileLength, totalLength - offset);
            Compute(offset, calcLength);
        }
    }
    __aicore__ inline void Compute(size_t offset, int calcLength)
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(calcLength * sizeof(T)), 
            0,
            0,
            0
        };
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x2Local, x2Gm[offset], copyParams, padParams);
        inQueueX2.EnQue(x2Local);
        LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
        DataCopyPad(x1Local, x1Gm[offset], copyParams, padParams);
        inQueueX1.EnQue(x1Local);

        LocalTensor<T> yLocal;
        x2Local = inQueueX2.DeQue<T>();
        // Power(x2Local, 2, x2Local);
        if constexpr (std::is_same_v<T, int32_t>) {
            ShiftLeft(x2Local, x2Local, (int32_t)23, calcLength);
            // Adds(x2Local, x2Local, (int32_t)0b00111111100000000000000000000000, calcLength);
            // Cast(x2Local, x2Local.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);
            x1Local = inQueueX1.DeQue<T>();
            Cast(x1Local.template ReinterpretCast<float>(), x1Local, RoundMode::CAST_TRUNC, calcLength);
            Add(x1Local, x1Local, x2Local, calcLength);
            inQueueX2.FreeTensor(x2Local);
            yLocal = outQueueY.AllocTensor<T>();
            Cast(yLocal, x1Local.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);
            // Mul(yLocal, x1Local, x2Local, calcLength);
            inQueueX1.FreeTensor(x1Local);
            // inQueueX2.FreeTensor(x2Local);
            outQueueY.EnQue(yLocal);
        }
        if constexpr (std::is_same_v<T, int16_t>) {
            ShiftLeft(x2Local, x2Local, (int16_t)10, calcLength);
            Adds(x2Local, x2Local, (int16_t)0b0011110000000000, calcLength);
            Cast(x2Local, x2Local.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            x1Local = inQueueX1.DeQue<T>();
            // Cast(x1Local.template ReinterpretCast<half>(), x1Local, RoundMode::CAST_TRUNC, calcLength);
            // Add(x1Local, x1Local, x2Local, calcLength);
            // inQueueX2.FreeTensor(x2Local);
            yLocal = outQueueY.AllocTensor<T>();
            // Cast(yLocal, x1Local.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            Mul(yLocal, x1Local, x2Local, calcLength);
            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);
            outQueueY.EnQue(yLocal);
        }
        if constexpr (std::is_same_v<T, int8_t>) {
            LocalTensor<int16_t> x2Local16 = calcQueueX2.AllocTensor<int16_t>();
            Cast(x2Local16.template ReinterpretCast<half>(), x2Local, RoundMode::CAST_NONE, calcLength);
            inQueueX2.FreeTensor(x2Local);
            Cast(x2Local16, x2Local16.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            // Adds(x2Local16, x2Local16, (int16_t)15, calcLength);
            ShiftLeft(x2Local16, x2Local16, (int16_t)10, calcLength);
            // Adds(x2Local16, x2Local16, (int16_t)0b0011110000000000, calcLength);
            // Cast(x2Local16, x2Local16.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            x1Local = inQueueX1.DeQue<T>();
            LocalTensor<int16_t> x1Local16 = calcQueueX1.AllocTensor<int16_t>();
            Cast(x1Local16.template ReinterpretCast<half>(), x1Local, RoundMode::CAST_NONE, calcLength);
            inQueueX1.FreeTensor(x1Local);
            // Cast(x1Local16, x1Local16.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            Add(x1Local16, x1Local16, x2Local16, calcLength);
            // Mul(yLocal16, x1Local16, x2Local16, calcLength);
            calcQueueX2.FreeTensor(x2Local16);
            yLocal = outQueueY.AllocTensor<T>();
            Cast(yLocal, x1Local16.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            calcQueueX1.FreeTensor(x1Local16);
            outQueueY.EnQue(yLocal);
        }

        yLocal = outQueueY.DeQue<T>();
        DataCopyPad(yGm[offset], yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }
};


template<typename T> class BitWiseLeftShift_64 {
private:
    TPipe* pipe;
    TQue<TPosition::VECIN, TQUE_NUM> inQueueX1, inQueueX2;
    TQue<TPosition::VECOUT, TQUE_NUM> outQueueY;
    TQue<TPosition::VECCALC, TQUE_NUM> calcQueueX1, calcQueueX2, calcQueueY;
    GlobalTensor<T> x1Gm, x2Gm, yGm;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;
    int blockNum;
    int blockIdx;
    int blockLength;
    int tailBlockNum;
    int tileLength;
    int tileNum;
    int bcLength;
    int lastLength;
    int lastLengthAlign;
    int blockLengthAlign;
    int tileLengthAlign;

public:
    __aicore__ inline BitWiseLeftShift_64() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        // this->totalLength1 = totalLength1;
        // this->totalLength2 = totalLength2;
        // this->totalLength = totalLength;
        // this->dimNum1 = dimNum1;
        // this->dimNum2 = dimNum2;
        // this->dimNum = 3;
        // this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < 3; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
            if (dim[i] != 1) {
                this->dimNum = i + 1;
            }
        }
        
        blockNum = GetBlockNum();
        blockIdx = GetBlockIdx();
        lastLength = dim[dimNum - 1];
        // 对齐到32字节
        lastLengthAlign = (lastLength * sizeof(int32_t) + 31) / 32 * 32 / sizeof(int32_t);
        bcLength = dim[dimNum - 2];
        blockLength = bcLength * lastLength;
        tileLength = lastLength;
        blockLengthAlign = bcLength * lastLengthAlign;
        tileLengthAlign = lastLengthAlign;
        int alreadyLength = blockIdx * blockLength;
        int alreadyLength2 = blockIdx * lastLength;

        // printf("blockIdx = %d, blockLength = %d, blockLengthAlign = %d\n", blockIdx, blockLength, blockLengthAlign);
        // printf("alreadyLength = %d, tileLength = %d, tileLengthAlign = %d\n", alreadyLength, tileLength, tileLengthAlign);

        x1Gm.SetGlobalBuffer((__gm__ T *)x1 + alreadyLength, blockLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2 + alreadyLength2, lastLength);
        yGm.SetGlobalBuffer((__gm__ T *)y + alreadyLength, blockLength);

        pipe->InitBuffer(inQueueX1, BUFFER_NUM_SP, tileLengthAlign * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM_SP, tileLengthAlign * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM_SP, tileLengthAlign * sizeof(T));
        pipe->InitBuffer(calcQueueX1, BUFFER_NUM_SP, tileLengthAlign * sizeof(int32_t));
        pipe->InitBuffer(calcQueueX2, BUFFER_NUM_SP, tileLengthAlign * sizeof(int32_t));
        pipe->InitBuffer(calcQueueY, BUFFER_NUM_SP, tileLengthAlign * sizeof(int32_t));
    }
    __aicore__ inline void Process()
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(lastLength * sizeof(T)), 
            0,
            0,
            0
        };
        int calcLength = tileLength;
        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x2Local, x2Gm, copyParams, padParams);
        inQueueX2.EnQue(x2Local);
        x2Local = inQueueX2.DeQue<T>();
        LocalTensor<int32_t> x2Local32 = calcQueueX2.AllocTensor<int32_t>();
        Cast(x2Local32, x2Local, RoundMode::CAST_NONE, lastLength);
        inQueueX2.FreeTensor(x2Local);
        ShiftLeft(x2Local32, x2Local32, (int32_t)23, calcLength);
        Adds(x2Local32, x2Local32, (int32_t)0b00111111100000000000000000000000, calcLength);
        Cast(x2Local32, x2Local32.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);

        for (int offset = 0; offset < blockLength; offset += lastLength) {
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm[offset], copyParams, padParams);
            inQueueX1.EnQue(x1Local);

            // x1Local = inQueueX1.DeQue<T>();
            // LocalTensor<int32_t> x2LocalBC = calcQueueX2BC.AllocTensor<int32_t>();
            // uint32_t dstShape[2] = {(uint32_t)bcLength, (uint32_t)lastLengthAlign};
            // uint32_t srcShape[2] = {(uint32_t)1, (uint32_t)lastLengthAlign};
            // Broadcast<float, 2, 0>(x2LocalBC.template ReinterpretCast<float>(), 
            //                         x2Local32.template ReinterpretCast<float>(), 
            //                         dstShape, 
            //                         srcShape);
            // Power(x2Local, 2, x2Local);
            // ShiftLeft(x2LocalBC, x2LocalBC, (int32_t)23, calcLength);
            // Adds(x2Local, x2Local, (int32_t)0b00111111100000000000000000000000, calcLength);
            // Cast(x2Local, x2Local.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);
            x1Local = inQueueX1.DeQue<T>();
            LocalTensor<int32_t> x1Local32 = calcQueueX1.AllocTensor<int32_t>();
            Cast(x1Local32, x1Local, RoundMode::CAST_NONE, calcLength);
            inQueueX1.FreeTensor(x1Local);
            // Cast(x1Local32.template ReinterpretCast<float>(), x1Local32, RoundMode::CAST_TRUNC, calcLength);
            // Add(x1Local32, x1Local32, x2LocalBC, calcLength);
            // calcQueueX2BC.FreeTensor(x2LocalBC);
            LocalTensor<int32_t> yLocal32 = calcQueueY.AllocTensor<int32_t>();
            Mul(yLocal32, x1Local32, x2Local32, calcLength);
            // Cast(yLocal32, x1Local32.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);
            calcQueueX1.FreeTensor(x1Local32);

            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
            // inQueueX1.FreeTensor(x1Local);
            // inQueueX2.FreeTensor(x2Local);
            Cast(yLocal, yLocal32, RoundMode::CAST_NONE, calcLength);
            calcQueueY.FreeTensor(yLocal32);
            outQueueY.EnQue(yLocal);

            yLocal = outQueueY.DeQue<T>();
            // copyParams = {
            //     (uint16_t)bcLength, 
            //     (uint32_t)(lastLength * sizeof(T)), 
            //     (uint32_t)(((lastLengthAlign - lastLength) * sizeof(T) + 31) / 32),
            //     0,
            //     0
            // };
            DataCopyPad(yGm[offset], yLocal, copyParams);
            outQueueY.FreeTensor(yLocal);
        }

        calcQueueX2.FreeTensor(x2Local32);
    }
};


template<typename T> class BitWiseLeftShift_CS3 {
private:
    TPipe* pipe;
    TQue<TPosition::VECIN, TQUE_NUM> inQueueX1, inQueueX2;
    TQue<TPosition::VECOUT, TQUE_NUM> outQueueY;
    TQue<TPosition::VECCALC, TQUE_NUM> calcQueueX1, calcQueueX2, calcQueueY;
    GlobalTensor<T> x1Gm, x2Gm, yGm;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;
    int blockNum;
    int blockIdx;
    int blockLength;
    int tailBlockNum;
    int tileLength;
    int tileNum;
    int calcLength;
    int bcLength;
    int lastLength;
    int lastLengthAlign;
    int blockLengthAlign;
    int tileLengthAlign;
    int idx0;
    int idx1;

public:
    __aicore__ inline BitWiseLeftShift_CS3() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        // this->totalLength1 = totalLength1;
        // this->totalLength2 = totalLength2;
        // this->totalLength = totalLength;
        // this->dimNum1 = dimNum1;
        // this->dimNum2 = dimNum2;
        // this->dimNum = 3;
        // this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < 3; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
            if (dim[i] != 1) {
                this->dimNum = i + 1;
            }
        }
        
        blockNum = GetBlockNum();
        blockIdx = GetBlockIdx();
        calcLength = dim[2];
        lastLength = dim[2];
        tileLength = lastLength;
        bcLength = dim[1];
        // this->blockLength = blockLength;
        blockLength = (dim[0] * dim[1] + blockNum - 1) / blockNum;
        int alreadyLength = blockIdx * blockLength;
        idx0 = alreadyLength / dim[1];
        idx1 = alreadyLength % dim[1];

        // printf("blockIdx = %d, blockLength = %d, blockLengthAlign = %d\n", blockIdx, blockLength, blockLengthAlign);
        // printf("alreadyLength = %d, tileLength = %d, tileLengthAlign = %d\n", alreadyLength, tileLength, tileLengthAlign);

        x1Gm.SetGlobalBuffer((__gm__ T *)x1, totalLength);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2, totalLength2);
        yGm.SetGlobalBuffer((__gm__ T *)y, totalLength);

        pipe->InitBuffer(inQueueX1, BUFFER_NUM_CS3, tileLength * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM_CS3, tileLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM_CS3, tileLength * sizeof(T));
        pipe->InitBuffer(calcQueueX1, BUFFER_NUM_CS3, tileLength * sizeof(int32_t));
        pipe->InitBuffer(calcQueueX2, BUFFER_NUM_CS3, tileLength * sizeof(int32_t));
        pipe->InitBuffer(calcQueueY, BUFFER_NUM_CS3, tileLength * sizeof(int32_t));
    }
    __aicore__ inline void Process()
    {
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(lastLength * sizeof(T)), 
            0,
            0,
            0
        };

        LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
        DataCopyPad(x2Local, x2Gm[idx0 * lastLength], copyParams, padParams);
        inQueueX2.EnQue(x2Local);
        x2Local = inQueueX2.DeQue<T>();
        LocalTensor<int32_t> x2Local32 = calcQueueX2.AllocTensor<int32_t>();
        Cast(x2Local32, x2Local, RoundMode::CAST_NONE, lastLength);
        inQueueX2.FreeTensor(x2Local);
        ShiftLeft(x2Local32, x2Local32, (int32_t)23, calcLength);
        Adds(x2Local32, x2Local32, (int32_t)0b00111111100000000000000000000000, calcLength);
        Cast(x2Local32, x2Local32.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);

        int offset1 = idx0 * bcLength * lastLength + idx1 * lastLength;
        for (int i = 0; i < blockLength && offset1 < totalLength; i++, offset1 += lastLength) {
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm[offset1], copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            x1Local = inQueueX1.DeQue<T>();
            LocalTensor<int32_t> x1Local32 = calcQueueX1.AllocTensor<int32_t>();
            Cast(x1Local32, x1Local, RoundMode::CAST_NONE, calcLength);
            inQueueX1.FreeTensor(x1Local);
            LocalTensor<int32_t> yLocal32 = calcQueueY.AllocTensor<int32_t>();
            Mul(yLocal32, x1Local32, x2Local32, calcLength);
            calcQueueX1.FreeTensor(x1Local32);
            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
            Cast(yLocal, yLocal32, RoundMode::CAST_NONE, calcLength);
            calcQueueY.FreeTensor(yLocal32);
            outQueueY.EnQue(yLocal);
            yLocal = outQueueY.DeQue<T>();
            DataCopyPad(yGm[offset1], yLocal, copyParams);
            outQueueY.FreeTensor(yLocal);

            idx1++;
            if (idx1 == dim[1]) {
                idx0++;
                idx1 = 0;

                calcQueueX2.FreeTensor(x2Local32);

                x2Local = inQueueX2.AllocTensor<T>();
                DataCopyPad(x2Local, x2Gm[idx0 * lastLength], copyParams, padParams);
                inQueueX2.EnQue(x2Local);
                x2Local = inQueueX2.DeQue<T>();
                x2Local32 = calcQueueX2.AllocTensor<int32_t>();
                Cast(x2Local32, x2Local, RoundMode::CAST_NONE, lastLength);
                inQueueX2.FreeTensor(x2Local);
                ShiftLeft(x2Local32, x2Local32, (int32_t)23, calcLength);
                Adds(x2Local32, x2Local32, (int32_t)0b00111111100000000000000000000000, calcLength);
                Cast(x2Local32, x2Local32.template ReinterpretCast<float>(), RoundMode::CAST_TRUNC, calcLength);
            }
        }
        calcQueueX2.FreeTensor(x2Local32);
    }
};



template<typename T> class BitWiseLeftShift_SP_BC {
private:
    TPipe* pipe;
    TQue<TPosition::VECIN, TQUE_NUM> inQueueX1, inQueueX2;
    TQue<TPosition::VECOUT, TQUE_NUM> outQueueY;
    GlobalTensor<T> x1Gm, x2Gm, yGm;
    int indices[5];
    ll totalLength1;
    ll totalLength2;
    ll totalLength;
    int dimNum1;
    int dimNum2;
    int dimNum;
    int* dim1;
    int* dim2;
    int* dim;
    int blockIdx;
    int blockLength;
    int blockLengthBC;
    int tailBlockNum;
    int tileLength;
    int tileLengthBC;
    int tileLengthBCAlign;
    int tileNum;
    int calcLength;
    int calcLengthBC;
    int bcLength;
    int bcLengthAlign;
    int alreadyLength;
    int alreadyLengthBC;

public:
    __aicore__ inline BitWiseLeftShift_SP_BC() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, TPipe* pipeIn, int* dim2, int* dim)
    {
        pipe = pipeIn;
        // this->totalLength1 = totalLength1;
        // this->totalLength2 = totalLength2;
        // this->totalLength = totalLength;
        // this->dimNum1 = dimNum1;
        // this->dimNum2 = dimNum2;
        // this->dimNum = 3;
        // this->dim1 = dim1;
        this->dim2 = dim2;
        this->dim = dim;

        this->totalLength = 1;
        this->totalLength2 = 1;
        for (int i = 0; i < 3; i++) {
            this->totalLength *= dim[i];
            this->totalLength2 *= dim2[i];
            if (dim[i] != 1) {
                this->dimNum = i + 1;
            }
        }
        
        // 最后一维广播的长度
        bcLength = dim[dimNum - 1];
        // 考虑32字节对齐
        bcLengthAlign = (bcLength * sizeof(T) + 32 - 1) / 32 * 32 / sizeof(T);
        int blockNum = GetBlockNum();
        blockIdx = GetBlockIdx();
        // 先给每个核分配要计算的长度
        // blockLength = totalLength2 / blockNum;
        // int tailBlockNum = totalLength2 % blockNum;
        // alreadyLength = blockLength * blockIdx;
        // 前面一些核要计算的长度多1
        // if (blockIdx < tailBlockNum) {
        //     alreadyLength += blockIdx;
        //     blockLength++;
        // } else {
        //     alreadyLength += tailBlockNum;
        // }
        blockLength = (totalLength2 + blockNum - 1) / blockNum;
        alreadyLength = blockLength * blockIdx;
        if (blockLength + alreadyLength > totalLength2) {
            blockLength = totalLength2 - alreadyLength;
        }
        blockLengthBC = blockLength * bcLength;
        alreadyLengthBC = alreadyLength * bcLength;
        // 每个循环里计算的长度
        tileLengthBC = 64000 / sizeof(T);
        if (tileLengthBC > blockLengthBC) {
            tileLengthBC = blockLengthBC;
        }
        tileLength = tileLengthBC / bcLength;
        tileLengthBC = tileLength * bcLength;
        // 考虑32字节对齐
        tileLengthBCAlign = tileLength * bcLengthAlign;
        // 循环次数
        tileNum = (blockLength + tileLength - 1) / tileLength;

        // printf("blockNum = %d, blockIdx = %d, blockLength = %d, blockLengthBC = %d\n", blockNum, blockIdx, blockLength, blockLengthBC);
        // printf("alreadyLength = %d, alreadyLengthBC = %d, tileLength = %d, tileLengthBC = %d, tileNum = %d\n", alreadyLength, alreadyLengthBC, tileLength, tileLengthBC, tileNum);

        x1Gm.SetGlobalBuffer((__gm__ T *)x1 + alreadyLengthBC, blockLengthBC);
        x2Gm.SetGlobalBuffer((__gm__ T *)x2 + alreadyLength, blockLength);
        yGm.SetGlobalBuffer((__gm__ T *)y + alreadyLengthBC, blockLengthBC);

        pipe->InitBuffer(inQueueX1, BUFFER_NUM_SP, tileLengthBCAlign * sizeof(T));
        pipe->InitBuffer(inQueueX2, BUFFER_NUM_SP, tileLength * sizeof(T));
        pipe->InitBuffer(outQueueY, BUFFER_NUM_SP, tileLengthBCAlign * sizeof(T));
    }
    __aicore__ inline void Process()
    {
        if (alreadyLength > totalLength2) {
            return;
        }
        calcLength = tileLength;
        calcLengthBC = tileLengthBC;
        DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
        DataCopyExtParams copyParamsBC = {
            (uint16_t)calcLength, 
            (uint32_t)(bcLength * sizeof(T)), 
            0,
            0,
            0
        };
        DataCopyExtParams copyParams = {
            (uint16_t)1, 
            (uint32_t)(calcLength * sizeof(T)), 
            0,
            0,
            0
        };
        int blockAlreadyLength = 0;
        int blockAlreadyLengthBC = 0;
        for (int tileIdx = 0; tileIdx < tileNum - 1; tileIdx++) {
            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm[blockAlreadyLengthBC], copyParamsBC, padParams);
            DataCopyPad(x2Local, x2Gm[blockAlreadyLength], copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            inQueueX2.EnQue(x2Local);

            x1Local = inQueueX1.DeQue<T>();
            x2Local = inQueueX2.DeQue<T>();
            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
            // Power(x2Local, 2, x2Local);

            // assert(std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>);

            // ShiftLeft(x2Local, x2Local, (int16_t)10, calcLength);
            // Adds(x2Local, x2Local, (int16_t)0b0011110000000000, calcLength);
            // Cast(x2Local, x2Local.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            // Mul(yLocal, x1Local, x2Local, calcLength);

            // 这次要计算calcLength行，每行有bcLength个元素
            int tileAlreadyLengthBCAlign = 0;
            for (int i = 0; i < calcLength; i++) {
                T x2 = x2Local(i);
                if constexpr (std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>) {
                    ShiftLeft(yLocal[tileAlreadyLengthBCAlign], x1Local[tileAlreadyLengthBCAlign], x2, bcLength);
                } else {
                    for (int j = 0; j < bcLength; j++) {
                        T x1 = x1Local(tileAlreadyLengthBCAlign + j);
                        T y = x1 << x2;
                        yLocal(tileAlreadyLengthBCAlign + j) = y;
                    }
                }
                // 考虑32字节对齐
                tileAlreadyLengthBCAlign += bcLengthAlign;
            }

            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);
            outQueueY.EnQue(yLocal);

            yLocal = outQueueY.DeQue<T>();
            DataCopyPad(yGm[blockAlreadyLengthBC], yLocal, copyParamsBC);
            outQueueY.FreeTensor(yLocal);

            blockAlreadyLength += calcLength;
            blockAlreadyLengthBC += calcLengthBC;
        }
        if (tileNum - 1 >= 0) {
            calcLength = blockLength - (tileNum - 1) * calcLength;
            calcLengthBC = calcLength * bcLengthAlign;
            copyParamsBC = {
                (uint16_t)calcLength, 
                (uint32_t)(bcLength * sizeof(T)), 
                0,
                0,
                0
            };
            copyParams = {
                (uint16_t)1, 
                (uint32_t)(calcLength * sizeof(T)), 
                0,
                0,
                0
            };

            LocalTensor<T> x1Local = inQueueX1.AllocTensor<T>();
            LocalTensor<T> x2Local = inQueueX2.AllocTensor<T>();
            DataCopyPad(x1Local, x1Gm[blockAlreadyLengthBC], copyParamsBC, padParams);
            DataCopyPad(x2Local, x2Gm[blockAlreadyLength], copyParams, padParams);
            inQueueX1.EnQue(x1Local);
            inQueueX2.EnQue(x2Local);

            x1Local = inQueueX1.DeQue<T>();
            x2Local = inQueueX2.DeQue<T>();
            LocalTensor<T> yLocal = outQueueY.AllocTensor<T>();
            // Power(x2Local, 2, x2Local);

            // assert(std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>);

            // ShiftLeft(x2Local, x2Local, (int16_t)10, calcLength);
            // Adds(x2Local, x2Local, (int16_t)0b0011110000000000, calcLength);
            // Cast(x2Local, x2Local.template ReinterpretCast<half>(), RoundMode::CAST_TRUNC, calcLength);
            // Mul(yLocal, x1Local, x2Local, calcLength);

            // 这次要计算calcLength行，每行有bcLength个元素
            int tileAlreadyLengthBCAlign = 0;
            for (int i = 0; i < calcLength; i++) {
                T x2 = x2Local(i);
                if constexpr (std::is_same_v<T, int16_t> || std::is_same_v<T, int32_t>) {
                    ShiftLeft(yLocal[tileAlreadyLengthBCAlign], x1Local[tileAlreadyLengthBCAlign], x2, bcLength);
                } else {
                    for (int j = 0; j < bcLength; j++) {
                        T x1 = x1Local(tileAlreadyLengthBCAlign + j);
                        T y = x1 << x2;
                        yLocal(tileAlreadyLengthBCAlign + j) = y;
                    }
                }
                // 考虑32字节对齐
                tileAlreadyLengthBCAlign += bcLengthAlign;
            }

            inQueueX1.FreeTensor(x1Local);
            inQueueX2.FreeTensor(x2Local);
            outQueueY.EnQue(yLocal);

            yLocal = outQueueY.DeQue<T>();
            DataCopyPad(yGm[blockAlreadyLengthBC], yLocal, copyParamsBC);
            outQueueY.FreeTensor(yLocal);
        }
    }
};


extern "C" __global__ __aicore__ void bitwise_left_shift(GM_ADDR input, GM_ADDR other, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    TPipe pipe;
    if (TILING_KEY_IS(1)) {
        BitWiseLeftShift_SP<DTYPE_INPUT> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    } else if (TILING_KEY_IS(3)) {
        BitWiseLeftShift_CS3<int64_t> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    } else if (TILING_KEY_IS(4)) {
        BitWiseLeftShift_SP_BC<DTYPE_INPUT> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    } else if (TILING_KEY_IS(5)) {
        BitWiseLeftShift<DTYPE_INPUT> op;
        op.Init(input, other, out, workspace, &pipe, tiling_data.dim2, tiling_data.dim);
        op.Process();
    }
}